#include "matrix_mem_view.hpp"

MatrixMemView::MatrixMemView(void* init_pos, int32_t row_count, int32_t column_count, 
                           size_t stride_0, size_t stride_1) {
    m_initial_pos = reinterpret_cast<size_t>(init_pos);
    m_current_pos = m_initial_pos;
    m_row_count = row_count;
    m_column_count = column_count;
    m_stride_0 = stride_0;
    m_stride_1 = stride_1;
    m_current_row = 0;
    m_current_column = 0;
}

void MatrixMemView::setCoordinate(int32_t column, int32_t row) {
    m_current_column = column;
    m_current_row = row;
    m_current_pos = m_initial_pos + m_current_column * m_stride_1 + m_current_row * m_stride_0;
}

void* MatrixMemView::next(size_t step) {
    size_t temp = m_current_row + step;
    if (temp < m_row_count) {
        m_current_pos += m_stride_0 * step;
        m_current_row = temp;
    } else {
        m_current_column += temp / m_row_count;
        m_current_row = temp % m_row_count;
        m_current_pos = m_initial_pos + m_current_column * m_stride_1 + m_current_row * m_stride_0;
    }
    return reinterpret_cast<void*>(m_current_pos);
}

void MatrixMemView::moveToNextColumn(bool keep_row) {
    m_current_column++;
    if (keep_row) {
        m_current_pos += m_stride_1;
    } else {
        m_current_pos = m_initial_pos + m_current_column * m_stride_1;
    }
}

void MatrixMemView::copyContinuousColumns(MatrixMemView& target, size_t step_length) {
    int32_t first_segment_length = m_row_count - m_current_row;
    size_t column_size = 0;
    
    if (first_segment_length > step_length) {
        column_size = step_length * m_stride_0;
        std::memcpy(target.getPointer(), getPointer(), column_size);
        m_current_pos += column_size;
        m_current_row += step_length;
        target.m_current_pos += column_size;
        target.m_current_row += step_length;
        return;
    }
    
    size_t remaining_length = step_length - first_segment_length;
    int32_t middle_columns_count = remaining_length / m_row_count;
    int32_t last_segment_length = remaining_length % m_row_count;
    
    column_size = first_segment_length * m_stride_0;
    std::memcpy(target.getPointer(), getPointer(), column_size);
    moveToNextColumn();
    target.moveToNextColumn();
    
    column_size = m_row_count * m_stride_0;
    for (int32_t i = 1; i < middle_columns_count + 1; i++) {
        std::memcpy(target.getPointer(), getPointer(), column_size);
        moveToNextColumn(true);
        target.moveToNextColumn(true);
    }
    
    if (last_segment_length) {
        column_size = last_segment_length * m_stride_0;
        std::memcpy(target.getPointer(), getPointer(), column_size);
        m_current_pos += column_size;
        target.m_current_pos += column_size;
    }
    
    m_current_row = last_segment_length;
}

void MatrixMemView::reinitialize() {
    m_current_row = 0;
    m_current_column = 0;
    m_current_pos = m_initial_pos;
}

char MatrixMemView::checkCompatibility(const MatrixMemView& target) const {
    char result_code = 0;
    if (target.m_column_count != m_column_count || target.m_row_count != m_row_count) {
        return 1;
    }
    
    const_cast<MatrixMemView&>(target).setCoordinate(m_current_column, m_current_row);
    
    if (target.m_stride_0 != m_stride_0) {
        result_code |= 2;
    }
    if (target.m_stride_1 != m_stride_1) {
        result_code |= 4;
    }
    return result_code;
}

void MatrixMemView::copyTo(MatrixMemView& target, size_t length, int32_t data_size) {
    void* destination = target.getPointer();
    void* source = getPointer();
    
    char compatibility_code = checkCompatibility(target);
    if (compatibility_code & 1) return;
    
    if (data_size == m_stride_0 && !(compatibility_code & 2)) {
        if (m_stride_1 == m_row_count * m_stride_0 && !(compatibility_code & 4)) {
            std::memcpy(destination, source, data_size * length);
            next(length);
            target.next(length);
        } else {
            copyContinuousColumns(target, length);
        }
    } else {
        for (size_t i = 0; i < length; i++) {
            std::memcpy(destination, source, data_size);
            destination = target.next();
            source = next();
        }
    }
}

template<typename DestinationType, typename SourceType>
void MatrixMemView::assignTo(MatrixMemView& target, size_t length) {
    void* destination = target.getPointer();
    void* source = getPointer();
    
    char compatibility_code = checkCompatibility(target);
    if (compatibility_code & 1) return;
    
    for (size_t i = 0; i < length; i++) {
        *(static_cast<DestinationType*>(destination)) = *(static_cast<SourceType*>(source));
        destination = target.next();
        source = next();
    }
}